#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <tuple>
#include <cmath>
#include <memory>
#include <Eigen/Dense>
#include <Eigen/Geometry>

using namespace Eigen;
using namespace std;

const double PI = 3.14159265358979323846;

class RobustKernel {
public:
    virtual ~RobustKernel() = default;
    virtual tuple<MatrixXd, VectorXd> robustify(const VectorXd& e) const = 0;
    virtual double weight(double squared_error) const = 0;
};

class HuberKernel : public RobustKernel {
public:
    HuberKernel(double delta = 1.0) : delta(delta), delta_sq(delta * delta) {}
    
    tuple<MatrixXd, VectorXd> robustify(const VectorXd& e) const override {
        double e_sq = e.squaredNorm();
        double w = weight(e_sq);
        VectorXd weighted_e = w * e;
        MatrixXd weighted_H = w * MatrixXd::Identity(e.size(), e.size());
        return make_tuple(weighted_H, weighted_e);
    }
    
    double weight(double squared_error) const override {
        return squared_error > delta_sq ? delta / sqrt(squared_error) : 1.0;
    }

private:
    double delta;
    double delta_sq;
};

struct Pose2D {
    double x, y, theta;
    Pose2D(double x = 0, double y = 0, double theta = 0) : x(x), y(y), theta(theta) {}
    
    Matrix3d toMatrix() const {
        Matrix3d T = Matrix3d::Identity();
        T(0, 0) = cos(theta); T(0, 1) = -sin(theta);
        T(1, 0) = sin(theta); T(1, 1) = cos(theta);
        T(0, 2) = x;
        T(1, 2) = y;
        return T;
    }
    
    static Pose2D fromMatrix(const Matrix3d& T) {
        return Pose2D(T(0, 2), T(1, 2), atan2(T(1, 0), T(0, 0)));
    }
};

struct PoseGraphEdge {
    int id1, id2;
    Pose2D measurement;
    Matrix3d information;
    bool is_outlier = false;
    
    PoseGraphEdge(int id1, int id2, const Pose2D& meas, const Matrix3d& info)
        : id1(id1), id2(id2), measurement(meas), information(info) {}
};

struct MarginalizationFactor {
    MatrixXd H_prior;
    VectorXd b_prior;
    unordered_set<int> related_vertices;
};

class SlidingWindowPoseGraph {
public:
    SlidingWindowPoseGraph(int window_size = 10) : window_size(window_size) {
        robust_kernel = make_shared<HuberKernel>(1.0);
    }
    
    void setRobustKernel(shared_ptr<RobustKernel> kernel) {
        robust_kernel = kernel;
    }
    
    int addVertex(const Pose2D& pose, bool is_fixed = false) {
        int id = next_vertex_id++;
        vertices[id] = pose;
        isFixed[id] = is_fixed;
        vertex_queue.push(id);
        
        if (vertex_queue.size() > window_size) {
            marginalizeOldestVertex();
        }
        
        return id;
    }
    
    void addEdge(const PoseGraphEdge& edge) {
        edges.push_back(edge);
        size_t edge_idx = edges.size() - 1;
        active_edges.insert(edge_idx);
        vertex_to_edges[edge.id1].insert(edge_idx);
        vertex_to_edges[edge.id2].insert(edge_idx);
    }
    
    double computeTotalError() {
        double total_error = 0.0;
        for (const auto& edge_idx : active_edges) {
            const auto& edge = edges[edge_idx];
            if (!vertices.count(edge.id1) || !vertices.count(edge.id2)) continue;
            
            Pose2D xi = vertices.at(edge.id1), xj = vertices.at(edge.id2);
            Matrix3d Tij = xi.toMatrix().inverse() * xj.toMatrix();
            Matrix3d T_error = edge.measurement.toMatrix().inverse() * Tij;
            Vector3d e(T_error(0, 2), T_error(1, 2), atan2(T_error(1, 0), T_error(0, 0)));
            total_error += e.transpose() * edge.information * e;
        }
        return total_error;
    }

    void optimize(int max_iterations = 10) {
        double mu = 1e-4;
        double v = 2.0;
        double min_error_reduction = 1e-6;
        double current_error = computeTotalError();

        for (int iter = 0; iter < max_iterations; ++iter) {
            buildSystem();
            auto backup_vertices = vertices;

            // LM阻尼
            MatrixXd H_lm = H;
            for (int i = 0; i < H_lm.rows(); ++i) {
                H_lm(i, i) += mu * H_lm(i, i);
            }

            // 求解
            VectorXd dx;
            LDLT<MatrixXd> ldlt(H_lm);
            if (ldlt.info() == Success) {
                dx = ldlt.solve(-b);
            } else {
                cerr << "Warning: LDLT failed, increasing damping." << endl;
                mu *= v;
                v *= 2.0;
                continue;
            }

            updatePoses(dx);
            double new_error = computeTotalError();
            double error_reduction = current_error - new_error;
            double predicted_reduction = 0.5 * (dx.dot(mu * H.diagonal().cwiseProduct(dx))) - 0.5 * b.dot(dx);

            double rho = (predicted_reduction > 0) ? (error_reduction / predicted_reduction) : 0;

            if (rho > 0.001) {
                current_error = new_error;
                mu *= max(1.0/3.0, 1.0 - pow(2.0*rho - 1.0, 3));
                v = 2.0;

                if (abs(error_reduction) < min_error_reduction) {
                    break;
                }
            } else {
                vertices = backup_vertices;
                mu *= v;
                v *= 2.0;
            }
        }
    }
    
    Pose2D getPose(int id) const {
        return vertices.at(id);
    }
    
    void printPoses() const {
        cout << "Current poses:" << endl;
        for (const auto& v : vertices) {
            cout << "Vertex " << v.first << ": (" << v.second.x << ", " 
                 << v.second.y << ", " << v.second.theta << ")" << endl;
        }
    }

private:
    int window_size;
    int next_vertex_id = 0;
    queue<int> vertex_queue;
    unordered_map<int, Pose2D> vertices;
    unordered_map<int, bool> isFixed;
    vector<PoseGraphEdge> edges;
    unordered_set<size_t> active_edges;
    unordered_map<int, unordered_set<size_t>> vertex_to_edges;
    unordered_map<int, MarginalizationFactor> marginalization_factors;
    shared_ptr<RobustKernel> robust_kernel;
    
    MatrixXd H;
    VectorXd b;
    unordered_map<int, int> var_indices;

    void buildSystem() {
        var_indices.clear();
        int free_vars = 0;
        for (const auto& v : vertices) {
            if (!isFixed[v.first]) {
                var_indices[v.first] = free_vars++;
            }
        }

        H = MatrixXd::Zero(3 * free_vars, 3 * free_vars);
        b = VectorXd::Zero(3 * free_vars);

        // 添加边约束
        for (size_t i = 0; i < edges.size(); ++i) {
            if (!active_edges.count(i)) continue;
            
            const auto& edge = edges[i];
            if (!vertices.count(edge.id1) || !vertices.count(edge.id2)) continue;

            Pose2D xi = vertices.at(edge.id1), xj = vertices.at(edge.id2);
            Matrix3d Tij = xi.toMatrix().inverse() * xj.toMatrix();
            Matrix3d T_error = edge.measurement.toMatrix().inverse() * Tij;
            
            Vector3d e;
            e << T_error(0, 2), T_error(1, 2), atan2(T_error(1, 0), T_error(0, 0));
            
            Matrix3d Ji = Matrix3d::Zero();
            Ji.block<2, 2>(0, 0) = -Rotation2Dd(-xi.theta).toRotationMatrix();
            Ji(2, 2) = -1;
            
            Matrix3d Jj = Matrix3d::Zero();
            Jj.block<2, 2>(0, 0) = Rotation2Dd(-xi.theta).toRotationMatrix();
            
            MatrixXd weighted_H;
            VectorXd weighted_e;
            std::tie(weighted_H, weighted_e) = robust_kernel->robustify(e);
            Matrix3d Omega_robust = edge.information * weighted_H;
            Vector3d e_robust = weighted_e;
            
            if (!isFixed[edge.id1] && !isFixed[edge.id2]) {
                int idx1 = var_indices[edge.id1], idx2 = var_indices[edge.id2];
                
                H.block<3, 3>(3*idx1, 3*idx1) += Ji.transpose() * Omega_robust * Ji;
                H.block<3, 3>(3*idx1, 3*idx2) += Ji.transpose() * Omega_robust * Jj;
                H.block<3, 3>(3*idx2, 3*idx1) += Jj.transpose() * Omega_robust * Ji;
                H.block<3, 3>(3*idx2, 3*idx2) += Jj.transpose() * Omega_robust * Jj;
                
                b.segment<3>(3*idx1) += Ji.transpose() * Omega_robust * e_robust;
                b.segment<3>(3*idx2) += Jj.transpose() * Omega_robust * e_robust;
            } else if (!isFixed[edge.id1]) {
                int idx1 = var_indices[edge.id1];
                H.block<3, 3>(3*idx1, 3*idx1) += Ji.transpose() * Omega_robust * Ji;
                b.segment<3>(3*idx1) += Ji.transpose() * Omega_robust * e_robust;
            } else if (!isFixed[edge.id2]) {
                int idx2 = var_indices[edge.id2];
                H.block<3, 3>(3*idx2, 3*idx2) += Jj.transpose() * Omega_robust * Jj;
                b.segment<3>(3*idx2) += Jj.transpose() * Omega_robust * e_robust;
            }
        }

        // 添加边缘化先验因子
        for (const auto& factor_pair : marginalization_factors) {
            const auto& factor = factor_pair.second;
            vector<int> active_vertices;
            
            for (int vid : factor.related_vertices) {
                if (vertices.count(vid) && !isFixed[vid]) {
                    active_vertices.push_back(vid);
                }
            }

            if (active_vertices.empty()) continue;

            MatrixXd H_local = MatrixXd::Zero(3 * active_vertices.size(), 
                                           3 * active_vertices.size());
            VectorXd b_local = VectorXd::Zero(3 * active_vertices.size());

            for (size_t i = 0; i < active_vertices.size(); ++i) {
                int global_i = var_indices[active_vertices[i]];
                for (size_t j = 0; j < active_vertices.size(); ++j) {
                    int global_j = var_indices[active_vertices[j]];
                    H_local.block<3,3>(3*i, 3*j) = 
                        factor.H_prior.block<3,3>(3*global_i, 3*global_j);
                }
                b_local.segment<3>(3*i) = factor.b_prior.segment<3>(3*global_i);
            }

            for (size_t i = 0; i < active_vertices.size(); ++i) {
                int global_i = var_indices[active_vertices[i]];
                for (size_t j = 0; j < active_vertices.size(); ++j) {
                    int global_j = var_indices[active_vertices[j]];
                    H.block<3,3>(3*global_i, 3*global_j) += 
                        H_local.block<3,3>(3*i, 3*j);
                }
                b.segment<3>(3*global_i) += b_local.segment<3>(3*i);
            }
        }
    }

    void updatePoses(const VectorXd& dx) {
        for (const auto& v : var_indices) {
            int idx = v.second;
            vertices[v.first].x += dx(3*idx);
            vertices[v.first].y += dx(3*idx + 1);
            vertices[v.first].theta = normalizeAngle(vertices[v.first].theta + dx(3*idx + 2));
        }
    }

    void marginalizeOldestVertex() {
        if (vertex_queue.empty()) return;
        int marg_id = vertex_queue.front();
        vertex_queue.pop();

        // 收集关联边和顶点
        unordered_set<size_t> related_edges = vertex_to_edges[marg_id];
        unordered_set<int> related_vertices;
        for (size_t edge_idx : related_edges) {
            const auto& edge = edges[edge_idx];
            if (edge.id1 != marg_id && vertices.count(edge.id1)) {
                related_vertices.insert(edge.id1);
            }
            if (edge.id2 != marg_id && vertices.count(edge.id2)) {
                related_vertices.insert(edge.id2);
            }
        }

        // 构建局部系统
        unordered_map<int, int> local_indices;
        int idx = 0;
        local_indices[marg_id] = idx++;
        for (int vid : related_vertices) {
            local_indices[vid] = idx++;
        }

        MatrixXd H_local = MatrixXd::Zero(3*idx, 3*idx);
        VectorXd b_local = VectorXd::Zero(3*idx);

        // 填充局部H和b
        for (size_t edge_idx : related_edges) {
            const auto& edge = edges[edge_idx];
            if (!vertices.count(edge.id1) || !vertices.count(edge.id2)) continue;
            
            Pose2D xi = vertices.at(edge.id1), xj = vertices.at(edge.id2);
            Matrix3d Tij = xi.toMatrix().inverse() * xj.toMatrix();
            Matrix3d T_error = edge.measurement.toMatrix().inverse() * Tij;
            
            Vector3d e;
            e << T_error(0, 2), T_error(1, 2), atan2(T_error(1, 0), T_error(0, 0));
            
            Matrix3d Ji = Matrix3d::Zero();
            Ji.block<2, 2>(0, 0) = -Rotation2Dd(-xi.theta).toRotationMatrix();
            Ji(2, 2) = -1;
            
            Matrix3d Jj = Matrix3d::Zero();
            Jj.block<2, 2>(0, 0) = Rotation2Dd(-xi.theta).toRotationMatrix();

            MatrixXd weighted_H;
            VectorXd weighted_e;
            std::tie(weighted_H, weighted_e) = robust_kernel->robustify(e);
            
            Matrix3d Omega_robust = edge.information * weighted_H;
            Vector3d e_robust = weighted_e;
            
            int idx1 = local_indices[edge.id1], idx2 = local_indices[edge.id2];
            
            H_local.block<3, 3>(3*idx1, 3*idx1) += Ji.transpose() * Omega_robust * Ji;
            H_local.block<3, 3>(3*idx1, 3*idx2) += Ji.transpose() * Omega_robust * Jj;
            H_local.block<3, 3>(3*idx2, 3*idx1) += Jj.transpose() * Omega_robust * Ji;
            H_local.block<3, 3>(3*idx2, 3*idx2) += Jj.transpose() * Omega_robust * Jj;
            
            b_local.segment<3>(3*idx1) += Ji.transpose() * Omega_robust * e_robust;
            b_local.segment<3>(3*idx2) += Jj.transpose() * Omega_robust * e_robust;
        }

        // 舒尔补计算
        if (related_vertices.size() > 0) {
            MatrixXd H_aa = H_local.block<3,3>(0, 0);
            MatrixXd H_ab = H_local.block(0, 3, 3, 3*(idx-1));
            MatrixXd H_ba = H_local.block(3, 0, 3*(idx-1), 3);
            MatrixXd H_bb = H_local.block(3, 3, 3*(idx-1), 3*(idx-1));

            VectorXd b_a = b_local.segment<3>(0);
            VectorXd b_b = b_local.segment(3, 3*(idx-1));

            // 确保H_aa可逆
            SelfAdjointEigenSolver<Matrix3d> eigensolver(H_aa);
            if (eigensolver.eigenvalues().minCoeff() < 1e-6) {
                H_aa += Matrix3d::Identity() * 1e-4;
            }

            MatrixXd H_aa_inv = H_aa.ldlt().solve(Matrix3d::Identity());
            MatrixXd H_prior = H_bb - H_ba * H_aa_inv * H_ab;
            VectorXd b_prior = b_b - H_ba * H_aa_inv * b_a;

            // 存储为先验因子
            MarginalizationFactor factor;
            factor.H_prior = H_prior;
            factor.b_prior = b_prior;
            factor.related_vertices = related_vertices;
            marginalization_factors[marg_id] = factor;
        }

        // 清理
        vertices.erase(marg_id);
        isFixed.erase(marg_id);
        for (size_t edge_idx : related_edges) {
            if (active_edges.count(edge_idx)) {
                const auto& edge = edges[edge_idx];
                vertex_to_edges[edge.id1].erase(edge_idx);
                vertex_to_edges[edge.id2].erase(edge_idx);
                active_edges.erase(edge_idx);
            }
        }
        vertex_to_edges.erase(marg_id);
    }

    static double normalizeAngle(double angle) {
        while (angle > PI) angle -= 2 * PI;
        while (angle < -PI) angle += 2 * PI;
        return angle;
    }
};

int main() {
    SlidingWindowPoseGraph optimizer(5);
    
    // 固定初始位姿
    int id0 = optimizer.addVertex(Pose2D(0, 0, 0), true);
    
    // 模拟机器人运动
    for (int i = 1; i <= 10; ++i) {
        Pose2D pose(i * 1.0, 0.1 * sin(i*0.5), 0.05 * cos(i*0.3));
        int id = optimizer.addVertex(pose);
        
        // 添加里程计边
        if (i > 1) {
            Matrix3d info = Matrix3d::Identity();
            info(2, 2) = 10.0;
            optimizer.addEdge(PoseGraphEdge(id-1, id, Pose2D(1.0, 0, 0), info));
        }
        
        // 添加回环
        if (i % 4 == 0 && i > 4) {
            Matrix3d loop_info = 100 * Matrix3d::Identity();
            optimizer.addEdge(PoseGraphEdge(id0, id, Pose2D(i*1.0, 0, 0), loop_info));
        }
        
        optimizer.optimize(5);
        cout << "After step " << i << ":" << endl;
        optimizer.printPoses();
        cout << "----------------------" << endl;
    }
    
    return 0;
}